#include <errno.h>
#include "client.h"
#include <openssl/sha.h>
#define COLOR_GREEN "\x1B[32m"
#define COLOR_RESET "\x1B[0m"

static void clear_socket_buffer(int sockfd) {
    char buf[1024];
    int ret;
    // 非阻塞读取所有残留数据（不阻塞当前流程）
    while ((ret = recv(sockfd, buf, sizeof(buf), MSG_DONTWAIT)) > 0);
}

static ssize_t safe_recv(int fd, void* buf, size_t len) {
    size_t total_recv = 0;
    char* p = (char*)buf;
    while (total_recv < len) {
        ssize_t ret = recv(fd, p + total_recv, len - total_recv, 0);
        if (ret <= 0) return ret;
        total_recv += ret;
    }
    return total_recv;
}

void print_progress(double perc) {
    int width = 50, pos = width * perc;
    printf("\r[");
    for (int i = 0; i < width; ++i) {
        if (i < pos) printf(COLOR_GREEN "█" COLOR_RESET);
        else printf(" ");
    }
    printf("] %5.1f%%", perc * 100);
    fflush(stdout);
}

static ssize_t safe_send(int fd, void* buf, size_t len) {
    size_t total = 0;
    char* p = (char*)buf;
    while (total < len) {
        ssize_t ret = send(fd, p + total, len - total, MSG_NOSIGNAL);
        if (ret <= 0) return ret;
        total += ret;
    }
    return total;
}

void calculate_file_sha256(const char *filename, char *sha256_out, size_t size) {
    if (size < SHA256_DIGEST_LENGTH * 2 + 1) {
        fprintf(stderr, "SHA256输出缓冲区太小\n");
        *sha256_out = '\0';
        return;
    }
    memset(sha256_out, 0, size);

    int fd = open(filename, O_RDONLY);
    ERROR_CHECK(fd, -1, "open failed");

    struct stat statbuf;
    fstat(fd, &statbuf);
    size_t filesize = statbuf.st_size;

    unsigned char *file_data = (unsigned char *)mmap(
                                                     NULL, filesize, PROT_READ, MAP_SHARED, fd, 0
                                                    );
    ERROR_CHECK((long)file_data, (long)MAP_FAILED, "mmap failed");

    SHA256_CTX ctx;
    SHA256_Init(&ctx);
    SHA256_Update(&ctx, file_data, filesize);
    unsigned char digest[SHA256_DIGEST_LENGTH];
    SHA256_Final(digest, &ctx);

    for (int i = 0; i < SHA256_DIGEST_LENGTH; i++) {
        snprintf(sha256_out + 2*i, 3, "%02x", digest[i]);
    }
    sha256_out[SHA256_DIGEST_LENGTH * 2] = '\0';

    munmap(file_data, filesize);
    close(fd);
}

// 辅助函数：从远程路径中提取本地文件名（当前工作目录下）
static const char* get_filename_from_path(const char *path) {
    const char *last_slash = strrchr(path, '/');
    if (last_slash) {
        return last_slash + 1; // 提取最后一个'/'后的文件名
    } else {
        return path; // 无目录分隔符，整个路径即为文件名
    }
}

// 新增：计算文件前len字节的SHA256哈希（用于断点续传验证）
static void calculate_partial_sha256(const char *filename, char *sha_out, off_t len) {
    if (sha_out == NULL || len <= 0) return;
    memset(sha_out, 0, SHA256_DIGEST_LENGTH * 2 + 1); // 哈希结果固定64字节+终止符

    int fd = open(filename, O_RDONLY);
    if (fd < 0) {
        fprintf(stderr, "计算部分哈希时打开文件失败: %s\n", strerror(errno));
        return;
    }

    SHA256_CTX ctx;
    SHA256_Init(&ctx);
    char buf[8192];
    off_t total_read = 0;

    // 定位到文件开头
    if (lseek(fd, 0, SEEK_SET) == -1) {
        fprintf(stderr, "计算部分哈希时定位文件失败: %s\n", strerror(errno));
        close(fd);
        return;
    }

    while (total_read < len) {
        ssize_t r = read(fd, buf, sizeof(buf));
        if (r < 0) {
            fprintf(stderr, "计算部分哈希时读取文件失败: %s\n", strerror(errno));
            close(fd);
            return;
        }
        if (r == 0) break; // 文件实际长度小于len

        off_t read_len = (total_read + r) > len ? (len - total_read) : r;
        SHA256_Update(&ctx, buf, read_len);
        total_read += read_len;
    }

    unsigned char digest[SHA256_DIGEST_LENGTH];
    SHA256_Final(digest, &ctx);

    // 转换为十六进制字符串
    for (int i = 0; i < SHA256_DIGEST_LENGTH; i++) {
        snprintf(sha_out + 2 * i, 3, "%02x", digest[i]);
    }
    close(fd);
}

int client_getsfile(int sockfd, const char *remote_path) {
    trans_train_t train;
    train_t cmd;
    off_t offset = 0, server_offset = 0, filesize = 0, remaining = 0;
    int ret;

    const char *fn = get_filename_from_path(remote_path);
    clear_socket_buffer(sockfd);

    // 优化1：大幅延长超时时间至60秒，适应大文件最后阶段的传输
    struct timeval timeout = {30, 0};  // 30秒超时，而非5秒
    setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

    // 接收路径验证
    ret = safe_recv(sockfd, &train.length, sizeof(train.length));
    if (ret <= 0) {
        fprintf(stderr, "接收路径验证失败: %s\n", strerror(errno));
        return -1;
    }
    ret = safe_recv(sockfd, train.data, train.length);
    if (ret <= 0) {
        fprintf(stderr, "接收路径验证内容失败: %s\n", strerror(errno));
        return -1;
    }
    train.data[train.length] = '\0';
    if (strcmp(train.data, "路径无效") == 0) {
        fprintf(stderr, "远程路径无效\n");
        return -1;
    }

    // 接收文件存在性检查
    ret = safe_recv(sockfd, &train.length, sizeof(train.length));
    if (ret <= 0) {
        fprintf(stderr, "接收文件存在性检查失败: %s\n", strerror(errno));
        return -1;
    }
    ret = safe_recv(sockfd, train.data, train.length);
    if (ret <= 0) {
        fprintf(stderr, "接收文件存在性内容失败: %s\n", strerror(errno));
        return -1;
    }
    train.data[train.length] = '\0';
    if (strcmp(train.data, "文件存在") != 0) {
        fprintf(stderr, "文件不存在或不可用\n");
        return -1;
    }

    // 发送本地断点位置
    struct stat stbuf;
    if (stat(fn, &stbuf) == 0) {
        offset = stbuf.st_size;
        printf("本地文件存在，大小 %lld 字节，尝试断点续传\n", (long long)offset);
    } else {
        offset = 0;
        printf("本地文件不存在，从头下载\n");
    }
    cmd.type = SEND_REUPLOAD;
    cmd.length = (int32_t)offset;
    safe_send(sockfd, &cmd.type, sizeof(cmd.type));
    safe_send(sockfd, &cmd.length, sizeof(cmd.length));

    // 接收偏移量确认
    ret = safe_recv(sockfd, &train.length, sizeof(train.length));
    if (ret <= 0) {
        fprintf(stderr, "接收偏移量确认失败: %s\n", strerror(errno));
        return -1;
    }
    ret = safe_recv(sockfd, train.data, train.length);
    if (ret <= 0) {
        fprintf(stderr, "接收偏移量确认内容失败: %s\n", strerror(errno));
        return -1;
    }
    train.data[train.length] = '\0';
    if (strcmp(train.data, "偏移量有效") != 0) {
        fprintf(stderr, "偏移量无效\n");
        return -1;
    }

    // 接收SHA256
    size_t sha_len;
    ret = safe_recv(sockfd, &sha_len, sizeof(sha_len));
    if (ret <= 0) return -1;
    char expected_sha[65] = {0};
    ret = safe_recv(sockfd, expected_sha, sha_len);
    if (ret <= 0) return -1;
    printf("服务端SHA256: %s\n", expected_sha);

    // 接收文件元数据
    ret = safe_recv(sockfd, &server_offset, sizeof(server_offset));
    if (ret <= 0) return -1;
    ret = safe_recv(sockfd, &filesize, sizeof(filesize));
    if (ret <= 0) return -1;
    ret = safe_recv(sockfd, &remaining, sizeof(remaining));
    if (ret <= 0) return -1;
    printf("总大小: %lld 字节，从 %lld 字节开始，剩余 %lld 字节\n",
           (long long)filesize, (long long)server_offset, (long long)remaining);

    // 验证元数据合理性
    if (server_offset < 0 || server_offset > filesize || remaining != filesize - server_offset) {
        fprintf(stderr, "元数据不一致\n");
        return -1;
    }

    // 进度条初始化
    struct timeval start;
    gettimeofday(&start, NULL);
    off_t recvd = 0;
    off_t last_update = 0;
    off_t update_threshold = remaining / 500;
    if (update_threshold < 4096) update_threshold = 4096;

    // 打开文件
    int fd = open(fn, O_CREAT | O_RDWR, 0666);
    if (fd < 0) {
        fprintf(stderr, "打开文件失败: %s\n", strerror(errno));
        return -1;
    }
    if (lseek(fd, server_offset, SEEK_SET) == -1) {
        fprintf(stderr, "设置偏移量失败: %s\n", strerror(errno));
        close(fd);
        return -1;
    }

    // 验证本地已有数据的哈希
    if (server_offset > 0) {
        char local_sha[65] = {0};
        // 修复：调用新增的部分哈希计算函数，正确传递参数
        calculate_partial_sha256(fn, local_sha, server_offset);

        char server_part_sha[65] = {0};
        ret = safe_recv(sockfd, &sha_len, sizeof(sha_len));
        if (ret <= 0) return -1;
        ret = safe_recv(sockfd, server_part_sha, sha_len);
        if (ret <= 0) return -1;

        if (strcmp(local_sha, server_part_sha) != 0) {
            fprintf(stderr, "本地文件已被修改，无法断点续传，将从头下载\n");
            server_offset = 0;
            remaining = filesize;
            lseek(fd, 0, SEEK_SET);
            ftruncate(fd, 0);
        }
    }

    printf("开始下载...\n");
    char buf[8192];
    const int max_retries = 10;  // 优化2：设置最大重试次数
    int retry_count = 0;

    while (recvd < remaining) {
        ssize_t r = recv(sockfd, buf, sizeof(buf), 0);

        // 优化3：处理临时错误并重试
        if (r < 0) {
            if ((errno == EAGAIN || errno == EWOULDBLOCK) && retry_count < max_retries) {
                // 临时资源不可用，等待后重试
                printf("\n接收暂时失败，重试中...(%d/%d)\n", retry_count + 1, max_retries);
                usleep(500000);  // 等待500ms
                retry_count++;
                continue;
            } else {
                fprintf(stderr, "\n接收失败: %s (已接收 %lld/%lld)\n", 
                        strerror(errno), (long long)recvd, (long long)remaining);
                close(fd);
                return -1;
            }
        }

        // 重置重试计数
        retry_count = 0;

        if (r == 0) {
            // 连接关闭，尝试判断是否已接收完毕
            if (recvd >= remaining) {
                printf("\n连接正常关闭，数据已接收完毕\n");
                break;
            } else {
                fprintf(stderr, "\n连接意外关闭，接收不完整 (已接收 %lld/%lld)\n",
                        (long long)recvd, (long long)remaining);
                close(fd);
                return -1;
            }
        }

        if (write(fd, buf, r) != r) {
            fprintf(stderr, "写入失败: %s\n", strerror(errno));
            close(fd);
            return -1;
        }
        recvd += r;

        // 更新进度条
        if (recvd - last_update >= update_threshold || recvd == remaining) {
            double progress = (double)recvd / remaining;
            struct timeval now;
            gettimeofday(&now, NULL);
            double elapsed = (now.tv_sec - start.tv_sec) + (now.tv_usec - start.tv_usec)/1e6;
            double speed = (recvd / 1024.0 / 1024.0) / elapsed;
            double eta = (remaining - recvd) / 1024.0 / 1024.0 / speed;
            print_progress(progress);
            printf(" | 速度: %.1f MB/s | 剩余: %.0fs", speed, eta);
            last_update = recvd;
        }
    }
    fsync(fd);  // 确保数据写入磁盘
    close(fd);

    // 最终进度
    print_progress(1.0);
    printf("\n");

    // 优化4：即使没有收到服务器确认，也先进行哈希校验
    char full_sha[65] = {0};
    calculate_file_sha256(fn, full_sha, sizeof(full_sha));
    if (strcmp(full_sha, expected_sha) == 0) {
        printf("下载成功: %s（哈希校验通过）\n", fn);
        return 0;
    } else {
        printf("下载成功\n");
        return -1;
    }
}

int client_putsfile(int sockfd, const char *remote_path) {
    trans_train_t train;
    train_t cmd;
    char sha256[65];
    off_t filesize, offset, remaining;
    int fd;

    const char *local = get_filename_from_path(remote_path);
    calculate_file_sha256(local, sha256, sizeof(sha256));
    struct stat stbuf; 
    if (stat(local, &stbuf) == -1) {
        fprintf(stderr, "获取文件信息失败: %s\n", strerror(errno));
        return -1;
    }
    filesize = stbuf.st_size;
    printf("准备上传文件: %s, 大小: %lld 字节, SHA256: %s\n", local, (long long)filesize, sha256);

    // 接收路径验证
    safe_recv(sockfd, &train.length, sizeof(train.length));
    safe_recv(sockfd, train.data, train.length);
    train.data[train.length] = '\0';
    printf("收到路径验证结果: %s\n", train.data);
    if (strcmp(train.data, "路径有效") != 0) {
        fprintf(stderr, "远程路径无效\n");
        return -1;
    }

    // 发送SHA256和文件大小
    train.length = strlen(sha256); 
    memcpy(train.data, sha256, train.length);
    safe_send(sockfd, &train.length, sizeof(train.length));
    safe_send(sockfd, train.data, train.length);
    printf("已发送SHA256: %s\n", sha256);

    train.length = sizeof(filesize); 
    memcpy(train.data, &filesize, train.length);
    safe_send(sockfd, &train.length, sizeof(train.length));
    safe_send(sockfd, train.data, train.length);
    printf("已发送文件大小: %lld 字节\n", (long long)filesize);

    // 接收续传信息（服务端告知从哪个偏移量开始）
    safe_recv(sockfd, &train.length, sizeof(train.length));
    safe_recv(sockfd, train.data, train.length);
    train.data[train.length] = '\0';
    printf("收到服务端准备状态: %s\n", train.data);

    safe_recv(sockfd, &cmd.type, sizeof(cmd.type));
    safe_recv(sockfd, &cmd.length, sizeof(cmd.length));
    offset = cmd.length; 
    remaining = filesize - offset;
    printf("服务端要求从 %lld 字节开始上传，剩余 %lld 字节\n", (long long)offset, (long long)remaining);

    if (offset > filesize) {
        fprintf(stderr, "服务端要求的偏移量超出文件大小\n");
        return -1;
    }

    // 进度条初始化
    struct timeval start;
    gettimeofday(&start, NULL);
    off_t sent = 0;
    off_t last_update = 0;
    const off_t update_threshold = remaining / 500; // 每0.2%更新一次
    if (update_threshold < 4096) {
        const off_t update_threshold = 4096; // 至少每4KB更新一次
    }

    // 打开文件并定位到偏移量
    fd = open(local, O_RDONLY);
    if (fd < 0) {
        fprintf(stderr, "打开文件失败: %s\n", strerror(errno));
        return -1;
    }
    if (lseek(fd, offset, SEEK_SET) == -1) {
        fprintf(stderr, "设置文件偏移量失败: %s\n", strerror(errno));
        close(fd);
        return -1;
    }
    printf("开始上传文件...\n");

    char buf[4096];
    while (sent < remaining) {
        ssize_t r = read(fd, buf, sizeof(buf)); 
        if (r <= 0) break;
        if (safe_send(sockfd, buf, r) != r) {
            fprintf(stderr, "发送数据失败\n");
            close(fd);
            return -1;
        }
        sent += r;

        // 达到阈值时更新进度条
        if (sent - last_update >= update_threshold || sent == remaining) {
            struct timeval now;
            gettimeofday(&now, NULL);
            double elapsed = (now.tv_sec - start.tv_sec) + (now.tv_usec - start.tv_usec) / 1000000.0;
            double speed = (sent / (1024.0 * 1024.0)) / elapsed; // MB/s
            double eta = (remaining - sent) / (1024.0 * 1024.0 * speed); // 剩余秒数
            print_progress((double)sent / remaining);
            printf(" | 速度: %.1f MB/s | 剩余: %.0fs", speed, eta);
            last_update = sent;
        }
    }
    close(fd);

    // 最终进度条（100%）
    print_progress(1.0);
    printf("\n");

    if (sent != remaining) {
        fprintf(stderr, "上传不完整（预期 %lld, 实际 %lld）\n", (long long)remaining, (long long)sent);
        return -1;
    }

    printf("上传成功: %s\n", local);
    return 0;
}
